"""Loader that loads image files."""
from typing import List, Callable, Optional

from langchain.document_loaders.unstructured import UnstructuredFileLoader
from langchain.text_splitter import TextSplitter
from langchain_core.documents import Document
from unstructured.partition.text import partition_text
import os
import fitz
from tqdm import tqdm
from typing import Union, Any, Dict
import numpy as np
import base64
import asyncio 


class UnstructuredPaddlePDFLoader(UnstructuredFileLoader):
    """Loader that uses unstructured to load image files, such as PNGs and JPGs."""
    def __init__(
        self,
        file_path: Union[str, List[str]],
        ocr_engine: Callable,
        mode: str = "single",
        **unstructured_kwargs: Any,
    ):
        """Initialize with file path."""
        self.ocr_engine = ocr_engine
        super().__init__(file_path=file_path, mode=mode, **unstructured_kwargs)

    def _get_elements(self) -> List:

        def create_tmp_file(filepath):
            full_dir_path = os.path.join(os.path.dirname(filepath), "temp_files")
            txt_file_path = os.path.join(full_dir_path, "{}.txt".format(os.path.split(filepath)[-1]))
            if not os.path.exists(full_dir_path):
                os.makedirs(full_dir_path)
            with open(txt_file_path, 'w', encoding='utf-8') as fout:
                doc_txt = self.ocr_engine(self.file_path)
                fout.write(doc_txt)

            return txt_file_path 
        

        txt_file_path = create_tmp_file(self.file_path)
        return partition_text(filename=txt_file_path, **self.unstructured_kwargs)
 